# flatten_squad.py
import json


def flatten_squad_json(input_file_path, output_file_path):
    """
    将SQuAD格式的JSON文件转换为扁平化的JSON Lines文件。
    每个问答对（Q&A pair）在新文件中占一行。
    """
    with open(input_file_path, 'r', encoding='utf-8') as f_in, \
            open(output_file_path, 'w', encoding='utf-8') as f_out:

        data = json.load(f_in)['data']

        for article in data:
            title = article.get('title', '')
            for paragraph in article['paragraphs']:
                context = paragraph['context']
                for qa in paragraph['qas']:
                    # 为每个问答对创建一个独立的、扁平的记录
                    flat_record = {
                        'id': qa['id'],
                        'title': title,
                        'context': context,
                        'question': qa['question'],
                        'answers': qa['answers']
                    }
                    # 将这个记录作为一行写入新文件
                    f_out.write(json.dumps(flat_record, ensure_ascii=False) + '\n')

    print(f"成功将 '{input_file_path}' 转换为扁平化的 '{output_file_path}'")


# --- 使用示例 ---
if __name__ == '__main__':
    # 将您之前的 train.json 和 eval.json 进行转换
    flatten_squad_json('train.json', 'train_flat.jsonl')
    flatten_squad_json('eval.json', 'eval_flat.jsonl')